package org.example;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.common.Role;
import org.example.utils.COMMON;

import java.util.ArrayList;
import java.util.List;
import java.util.Scanner;

public class MultiConversation {

    private static List<Message> messages = new ArrayList<>();

    private static Integer token = 0;

    public static GenerationParam createGenerationParam() {
        return GenerationParam.builder()
                .apiKey(COMMON.API_KEY)
                .model(COMMON.MODEL)
                .messages(messages)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .build();
    }
    public static GenerationResult callGenerationWithMessages(GenerationParam param) throws ApiException, NoApiKeyException, InputRequiredException {
        Generation gen = new Generation();
        return gen.call(param);
    }
    private static Message createMessage(Role role, String content) {
        return Message.builder().role(role.getValue()).content(content).build();
    }
    public static void main(String[] args) {
        try {
            Scanner scanner = new Scanner(System.in);
            Integer times = 1;
            System.out.println("欢迎来到对话！(若需退出请输入exit)");
            messages.add(createMessage(Role.SYSTEM, "you are a helpful assistant"));
            while(true){

                System.out.println("第"+times+"轮（当前消耗总token:"+ token +"）");
                System.out.print("用户：");
                String input = scanner.nextLine();
                if(input.equals("exit")){
                    break;
                }
                messages.add(createMessage(Role.USER, input));

                // 开始对话
                GenerationParam param = createGenerationParam();
                GenerationResult result = callGenerationWithMessages(param);

                // 获取AI回复的内容
                String answer = result.getOutput().getChoices().get(0).getMessage().getContent();
                messages.add(createMessage(Role.ASSISTANT, answer));   // 将AI的回复添加到messages中
                System.out.println("AI："+answer);

                // 计算已消耗token
                token += result.getUsage().getTotalTokens();
                times++;
            }
        } catch (Exception e) {
            System.err.println("错误信息：" + e.getMessage());
            System.out.println("请参考文档：https://help.aliyun.com/zh/model-studio/developer-reference/error-code");
        }
    }
}
